{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# DiscreteBayesNet"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "license_cell",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\nAtlanta, Georgia 30332-0415\nAll Rights Reserved\n\nAuthors: Frank Dellaert, et al. (see THANKS for the full author list)\n\nSee LICENSE for the license information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/gtsam/discrete/doc/DiscreteBayesNet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam\n",
        "except ImportError:\n",
        "    pass  # Not running on Colab, do nothing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A `DiscreteBayesNet` is a `BayesNet` of `DiscreteConditional`s. It represents a joint probability distribution over a set of discrete variables as a product of conditional probabilities:\n",
        "$$\n",
        "P(X_1, ..., X_n) = \\prod_i P(X_i | \\text{Parents}(X_i))\n",
        "$$\n",
        "Each conditional $P(X_i | \\text{Parents}(X_i))$ is a `gtsam.DiscreteConditional`.\n",
        "\n",
        "A `DiscreteBayesNet` is typically obtained by eliminating a `DiscreteFactorGraph`, but it can also be constructed manually."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "import gtsam\n",
        "import numpy as np\n",
        "import graphviz\n",
        "\n",
        "from gtsam.symbol_shorthand import C, S, R, W"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a DiscreteBayesNet\n",
        "\n",
        "We will manually construct the classic \"Wet Grass\" Bayes Net, which models the relationships between Cloudy, Sprinkler, Rain, and Wet Grass."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Manually Constructed DiscreteBayesNet:\n",
            "DiscreteBayesNet\n",
            " \n",
            "size: 4\n",
            "conditional 0:  P( w0 | s0 r0 ):\n",
            " Choice(w0) \n",
            " 0 Choice(s0) \n",
            " 0 0 Choice(r0) \n",
            " 0 0 0 Leaf    1\n",
            " 0 0 1 Leaf  0.1\n",
            " 0 1 Choice(r0) \n",
            " 0 1 0 Leaf  0.1\n",
            " 0 1 1 Leaf 0.01\n",
            " 1 Choice(s0) \n",
            " 1 0 Choice(r0) \n",
            " 1 0 0 Leaf    0\n",
            " 1 0 1 Leaf  0.9\n",
            " 1 1 Choice(r0) \n",
            " 1 1 0 Leaf  0.9\n",
            " 1 1 1 Leaf 0.99\n",
            "\n",
            "conditional 1:  P( s0 | c0 ):\n",
            " Choice(s0) \n",
            " 0 Choice(c0) \n",
            " 0 0 Leaf  0.5\n",
            " 0 1 Leaf  0.9\n",
            " 1 Choice(c0) \n",
            " 1 0 Leaf  0.5\n",
            " 1 1 Leaf  0.1\n",
            "\n",
            "conditional 2:  P( r0 | c0 ):\n",
            " Choice(r0) \n",
            " 0 Choice(c0) \n",
            " 0 0 Leaf  0.8\n",
            " 0 1 Leaf  0.2\n",
            " 1 Choice(c0) \n",
            " 1 0 Leaf  0.2\n",
            " 1 1 Leaf  0.8\n",
            "\n",
            "conditional 3:  P( c0 ):\n",
            " Leaf  0.5\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Define keys for the variables (all binary)\n",
        "Cloudy = (C(0), 2)\n",
        "Sprinkler = (S(0), 2)\n",
        "Rain = (R(0), 2)\n",
        "WetGrass = (W(0), 2)\n",
        "\n",
        "# Create an empty Bayes Net\n",
        "dbn = gtsam.DiscreteBayesNet()\n",
        "\n",
        "# P(WetGrass | Sprinkler, Rain)\n",
        "# S,R=FF: 1.0/0.0; S,R=FT: 0.1/0.9; S,R=TF: 0.1/0.9; S,R=TT: 0.01/0.99\n",
        "dbn.add(WetGrass, [Sprinkler, Rain], \"1.0/0.0 0.1/0.9 0.1/0.9 0.01/0.99\")\n",
        "\n",
        "# P(Sprinkler | Cloudy)\n",
        "# C=F: 0.5/0.5, C=T: 0.9/0.1\n",
        "dbn.add(Sprinkler, [Cloudy], \"0.5/0.5 0.9/0.1\")\n",
        "\n",
        "# P(Rain | Cloudy)\n",
        "# C=F: 0.8/0.2, C=T: 0.2/0.8\n",
        "dbn.add(Rain, [Cloudy], \"0.8/0.2 0.2/0.8\")\n",
        "\n",
        "# P(Cloudy) - a prior\n",
        "dbn.add(Cloudy, \"0.5/0.5\")\n",
        "\n",
        "print(\"Manually Constructed DiscreteBayesNet:\")\n",
        "dbn.print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/svg+xml": [
              "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
              "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
              " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
              "<!-- Generated by graphviz version 12.0.0 (0)\n",
              " -->\n",
              "<!-- Pages: 1 -->\n",
              "<svg width=\"134pt\" height=\"188pt\"\n",
              " viewBox=\"0.00 0.00 134.00 188.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
              "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 184)\">\n",
              "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-184 130,-184 130,4 -4,4\"/>\n",
              "<!-- var7133701809754865664 -->\n",
              "<g id=\"node1\" class=\"node\">\n",
              "<title>var7133701809754865664</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"63\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"63\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">c0</text>\n",
              "</g>\n",
              "<!-- var8214565720323784704 -->\n",
              "<g id=\"node2\" class=\"node\">\n",
              "<title>var8214565720323784704</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"27\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">r0</text>\n",
              "</g>\n",
              "<!-- var7133701809754865664&#45;&gt;var8214565720323784704 -->\n",
              "<g id=\"edge1\" class=\"edge\">\n",
              "<title>var7133701809754865664&#45;&gt;var8214565720323784704</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M54.65,-144.76C50.42,-136.55 45.19,-126.37 40.42,-117.09\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"43.68,-115.79 36,-108.49 37.46,-118.99 43.68,-115.79\"/>\n",
              "</g>\n",
              "<!-- var8286623314361712640 -->\n",
              "<g id=\"node3\" class=\"node\">\n",
              "<title>var8286623314361712640</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"99\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"99\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">s0</text>\n",
              "</g>\n",
              "<!-- var7133701809754865664&#45;&gt;var8286623314361712640 -->\n",
              "<g id=\"edge2\" class=\"edge\">\n",
              "<title>var7133701809754865664&#45;&gt;var8286623314361712640</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M71.35,-144.76C75.58,-136.55 80.81,-126.37 85.58,-117.09\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"88.54,-118.99 90,-108.49 82.32,-115.79 88.54,-118.99\"/>\n",
              "</g>\n",
              "<!-- var8574853690513424384 -->\n",
              "<g id=\"node4\" class=\"node\">\n",
              "<title>var8574853690513424384</title>\n",
              "<ellipse fill=\"none\" stroke=\"black\" cx=\"63\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
              "<text text-anchor=\"middle\" x=\"63\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">w0</text>\n",
              "</g>\n",
              "<!-- var8214565720323784704&#45;&gt;var8574853690513424384 -->\n",
              "<g id=\"edge4\" class=\"edge\">\n",
              "<title>var8214565720323784704&#45;&gt;var8574853690513424384</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M35.35,-72.76C39.58,-64.55 44.81,-54.37 49.58,-45.09\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"52.54,-46.99 54,-36.49 46.32,-43.79 52.54,-46.99\"/>\n",
              "</g>\n",
              "<!-- var8286623314361712640&#45;&gt;var8574853690513424384 -->\n",
              "<g id=\"edge3\" class=\"edge\">\n",
              "<title>var8286623314361712640&#45;&gt;var8574853690513424384</title>\n",
              "<path fill=\"none\" stroke=\"black\" d=\"M90.65,-72.76C86.42,-64.55 81.19,-54.37 76.42,-45.09\"/>\n",
              "<polygon fill=\"black\" stroke=\"black\" points=\"79.68,-43.79 72,-36.49 73.46,-46.99 79.68,-43.79\"/>\n",
              "</g>\n",
              "</g>\n",
              "</svg>\n"
            ],
            "text/plain": [
              "<graphviz.sources.Source at 0x10ef12060>"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Visualize the Bayes Net structure\n",
        "graphviz.Source(dbn.dot())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Operations on DiscreteBayesNet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Joint Probability P(C=T,S=F,R=T,W=T): 0.3240\n",
            "Log Probability: -1.1270\n",
            "\n",
            "Sampled values:\n",
            "DiscreteValues{7133701809754865664: 1, 8214565720323784704: 1, 8286623314361712640: 0, 8574853690513424384: 1}\n",
            "\n",
            "MAP Solution (Optimize):\n",
            "DiscreteValues{7133701809754865664: 1, 8214565720323784704: 1, 8286623314361712640: 0, 8574853690513424384: 1}\n"
          ]
        }
      ],
      "source": [
        "# --- Evaluation ---\n",
        "# Calculate the joint probability of a specific assignment of variables\n",
        "values = gtsam.DiscreteValues()\n",
        "values[C(0)] = 1 # Cloudy = T\n",
        "values[S(0)] = 0 # Sprinkler = F\n",
        "values[R(0)] = 1 # Rain = T\n",
        "values[W(0)] = 1 # WetGrass = T\n",
        "\n",
        "# P(C=T,S=F,R=T,W=T) = P(C=T)*P(S=F|C=T)*P(R=T|C=T)*P(W=T|S=F,R=T)\n",
        "#                     = 0.5 * 0.9 * 0.8 * 0.9 = 0.324\n",
        "\n",
        "prob = dbn.evaluate(values)\n",
        "log_prob = dbn.logProbability(values)\n",
        "print(f\"Joint Probability P(C=T,S=F,R=T,W=T): {prob:.4f}\")\n",
        "print(f\"Log Probability: {log_prob:.4f}\")\n",
        "\n",
        "# --- Sampling ---\n",
        "# Perform ancestral sampling to get a complete assignment.\n",
        "sample = dbn.sample()\n",
        "print(\"\\nSampled values:\")\n",
        "print(sample)\n",
        "\n",
        "# --- Optimization (Finding MPE state) ---\n",
        "# Finds the most probable explanation (MPE), the assignment with the highest probability.\n",
        "fg = gtsam.DiscreteFactorGraph(dbn)\n",
        "mpe_solution = fg.optimize()\n",
        "print(\"\\nMAP Solution (Optimize):\")\n",
        "print(mpe_solution)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Displaying with Markdown/HTML\n",
        "The rich display for a Bayes Net shows the CPT for each conditional."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div><p><tt>DiscreteBayesNet</tt> of size 4</p><div>\n",
              "<p>  <i>P(w0|s0,r0):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>s0</i></th><th><i>r0</i></th><th>0</th><th>1</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><th>0</th><td>1</td><td>0</td></tr>\n",
              "    <tr><th>0</th><th>1</th><td>0.1</td><td>0.9</td></tr>\n",
              "    <tr><th>1</th><th>0</th><td>0.1</td><td>0.9</td></tr>\n",
              "    <tr><th>1</th><th>1</th><td>0.01</td><td>0.99</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "<div>\n",
              "<p>  <i>P(s0|c0):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>c0</i></th><th>0</th><th>1</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><td>0.5</td><td>0.5</td></tr>\n",
              "    <tr><th>1</th><td>0.9</td><td>0.1</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "<div>\n",
              "<p>  <i>P(r0|c0):</i></p>\n",
              "<table class='DiscreteConditional'>\n",
              "  <thead>\n",
              "    <tr><th><i>c0</i></th><th>0</th><th>1</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><td>0.8</td><td>0.2</td></tr>\n",
              "    <tr><th>1</th><td>0.2</td><td>0.8</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "<div>\n",
              "<p>  <i>P(c0):</i></p>\n",
              "<div>\n",
              "<table class='DecisionTreeFactor'>\n",
              "  <thead>\n",
              "    <tr><th>c0</th><th>value</th></tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr><th>0</th><td>0.5</td></tr>\n",
              "    <tr><th>1</th><td>0.5</td></tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n"
            ],
            "text/markdown": [
              "`DiscreteBayesNet` of size 4\n",
              "\n",
              " *P(w0|s0,r0):*\n",
              "\n",
              "|*s0*|*r0*|0|1|\n",
              "|:-:|:-:|:-:|:-:|\n",
              "|0|0|1|0|\n",
              "|0|1|0.1|0.9|\n",
              "|1|0|0.1|0.9|\n",
              "|1|1|0.01|0.99|\n",
              "\n",
              " *P(s0|c0):*\n",
              "\n",
              "|*c0*|0|1|\n",
              "|:-:|:-:|:-:|\n",
              "|0|0.5|0.5|\n",
              "|1|0.9|0.1|\n",
              "\n",
              " *P(r0|c0):*\n",
              "\n",
              "|*c0*|0|1|\n",
              "|:-:|:-:|:-:|\n",
              "|0|0.8|0.2|\n",
              "|1|0.2|0.8|\n",
              "\n",
              " *P(c0):*\n",
              "\n",
              "|c0|value|\n",
              "|:-:|:-:|\n",
              "|0|0.5|\n",
              "|1|0.5|\n",
              "\n"
            ],
            "text/plain": [
              "DiscreteBayesNet\n",
              " \n",
              "size: 4\n",
              "conditional 0:  P( w0 | s0 r0 ):\n",
              " Choice(w0) \n",
              " 0 Choice(s0) \n",
              " 0 0 Choice(r0) \n",
              " 0 0 0 Leaf    1\n",
              " 0 0 1 Leaf  0.1\n",
              " 0 1 Choice(r0) \n",
              " 0 1 0 Leaf  0.1\n",
              " 0 1 1 Leaf 0.01\n",
              " 1 Choice(s0) \n",
              " 1 0 Choice(r0) \n",
              " 1 0 0 Leaf    0\n",
              " 1 0 1 Leaf  0.9\n",
              " 1 1 Choice(r0) \n",
              " 1 1 0 Leaf  0.9\n",
              " 1 1 1 Leaf 0.99\n",
              "\n",
              "conditional 1:  P( s0 | c0 ):\n",
              " Choice(s0) \n",
              " 0 Choice(c0) \n",
              " 0 0 Leaf  0.5\n",
              " 0 1 Leaf  0.9\n",
              " 1 Choice(c0) \n",
              " 1 0 Leaf  0.5\n",
              " 1 1 Leaf  0.1\n",
              "\n",
              "conditional 2:  P( r0 | c0 ):\n",
              " Choice(r0) \n",
              " 0 Choice(c0) \n",
              " 0 0 Leaf  0.8\n",
              " 0 1 Leaf  0.2\n",
              " 1 Choice(c0) \n",
              " 1 0 Leaf  0.2\n",
              " 1 1 Leaf  0.8\n",
              "\n",
              "conditional 3:  P( c0 ):\n",
              " Leaf  0.5\n"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dbn"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "py312",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}